import numpy as np
from typing import *
from utilities import *
import re

class Part:
    def __init__(self, name=None):
        self.name = name
        self.node_array = np.empty((0, 3), dtype=np.float32)
        self.elem_types:Dict[str, np.ndarray] = {}  # key=单元类型  value=相关节点编号的numpy数组
        self.node_id_to_index:Dict[int,int] = {}  # 节点编号转数组索引
        self.elem_info:Dict[int, tuple[str,int]] = {}  # 单元编号到(类型str, 数组索引)的映射
        self.elem_type_offsets:Dict[str, int] = {}

    def elem_summary(self):
        sum_ = 0
        for et,array in self.elem_types.items():
            self.elem_type_offsets[et] = sum_
            sum_ += len(array)

    def get_elem_seq_index(self, elem_id:int):
        et, idx = self.elem_info.get(elem_id, ("", -1))
        if idx == -1 : return -1
        return self.elem_type_offsets[et] + idx
    
    @property
    def elem_total_num(self):
        return sum((len(arr) for arr in self.elem_types.values()))
    
class Instance:
    def __init__(self, part:Part, name=None):
        self.ref_part = part
        self.name = name
        self.translate = np.zeros(3, dtype=np.float32)  # 平移向量
        self.rotation = None  # 旋转信息 (旋转轴点1, 旋转轴点2, 旋转角度(弧度))
        
class Assembly:
    def __init__(self):
        self.instances:List[Instance] = []
        self.parts = {}

class AbaqusInpParser:
    def __init__(self, filename):
        self.filename = filename
        self.assembly = Assembly()
        self.current_part = None
        self.current_instance = None
        self.current_elem_type = None
        
    def parse(self):
        with open(self.filename, 'r') as f:
            lines = f.readlines()
            
        i = 0
        while i < len(lines):
            line = lines[i].strip()
            
            # 跳过空行和注释
            if not line or line.startswith('**'):
                i += 1
                continue
                
            # 处理部件定义
            if line.startswith('*Part, name='):
                part_name = line.split('=')[1]
                self.current_part = Part(part_name)
                self.assembly.parts[part_name] = self.current_part
                i += 1
                
            # 处理部件结束
            elif line.startswith('*End Part'):
                self.current_part.elem_summary()
                self.current_part = None
                i += 1
                
            # 处理节点数据
            elif line.startswith('*Node'):
                i, node_data = self._parse_node_block(lines, i+1)
                if self.current_part:
                    self._process_node_data(node_data)
                else:i += 1
                
            # 处理单元数据
            elif line.startswith('*Element'):
                # 解析单元类型
                elem_type = self._parse_element_type(line)
                if self.current_part:
                    self.current_elem_type = elem_type
                    i, elem_data = self._parse_element_block(lines, i+1)
                    self._process_element_data(elem_data)
                else:i += 1
                
            # 处理实例定义
            elif line.startswith('*Instance, name='):
                instance_name = line.split('=')[1].split(',')[0]
                instance_name = instance_name.upper()
                part_name = None
                if 'part=' in line:
                    part_name = line.split('part=')[1].strip()
                
                if part_name and part_name in self.assembly.parts:
                    self.current_instance = Instance(self.assembly.parts[part_name], instance_name)
                    self.assembly.instances.append(self.current_instance)
                    
                    # 解析实例的平移和旋转数据
                    i += 1
                    if i < len(lines) and not lines[i].strip().startswith('*'):
                        # 第一行是平移
                        translate_line = lines[i].strip()
                        translate = self._parse_translate(translate_line)
                        if translate is not None:
                            self.current_instance.translate = translate
                        
                        # 检查是否有旋转数据
                        i += 1
                        if i < len(lines) and not lines[i].strip().startswith('*'):
                            rotate_line = lines[i].strip()
                            rotation = self._parse_rotate(rotate_line)
                            if rotation is not None:
                                self.current_instance.rotation = rotation
                            i += 1
                
            # 处理实例结束
            elif line.startswith('*End Instance'):
                self.current_instance = None
                i += 1
                
            else:
                i += 1
                
        return self.assembly
    
    def _parse_node_block(self, lines, start_idx):
        node_data = []
        i = start_idx
        while i < len(lines):
            line = lines[i].strip()
            if line.startswith('*'):
                break
            if line:
                parts = [p for p in re.split(r'[,\s]+', line) if p]
                try:
                    node_id = int(parts[0])
                    coords = list(map(float, parts[1:4]))
                    node_data.append((node_id, coords))
                except (ValueError, IndexError):
                    pass
            i += 1
        return i, node_data
    
    def _process_node_data(self, node_data):
        if not node_data:
            return
            
        # 按节点ID排序
        node_data.sort(key=lambda x: x[0])
        
        # 创建节点数组和映射
        num_nodes = len(node_data)
        self.current_part.node_array = np.empty((num_nodes, 3), dtype=np.float32)
        for idx, (node_id, coords) in enumerate(node_data):
            self.current_part.node_array[idx] = coords
            self.current_part.node_id_to_index[node_id] = idx
    
    def _parse_element_type(self, line):
        # 从元素定义行中提取元素类型
        if 'type=' in line:
            return line.split('type=')[1].strip()
        return 'UNKNOWN'
    
    def _parse_element_block(self, lines, start_idx):
        elem_data = []
        i = start_idx
        while i < len(lines):
            line = lines[i].strip()
            if line.startswith('*'):
                break
            if line:
                parts = [p for p in re.split(r'[,\s]+', line) if p]
                try:
                    elem_id = int(parts[0])
                    nodes = list(map(int, parts[1:]))
                    elem_data.append((elem_id, nodes))
                except (ValueError, IndexError):
                    pass
            i += 1
        return i, elem_data
    
    def _process_element_data(self, elem_data):
        if not elem_data:
            return
            
        # 按单元ID排序
        elem_data.sort(key=lambda x: x[0])
        
        # 确定单元节点数
        elem_nodes = len(elem_data[0][1])
            
        # 创建单元数组和映射
        num_elems = len(elem_data)
        if self.current_elem_type not in self.current_part.elem_types:
            self.current_part.elem_types[self.current_elem_type] = np.empty((0, elem_nodes), dtype=np.int32)
            
        existing_array = self.current_part.elem_types[self.current_elem_type]
        new_array = np.array([nodes for _, nodes in elem_data], dtype=np.int32)
        self.current_part.elem_types[self.current_elem_type] = np.vstack((existing_array, new_array))
        
        # 更新elem_info (单元编号 -> (类型, 数组索引))
        offset = existing_array.shape[0]
        for idx, (elem_id, _) in enumerate(elem_data):
            self.current_part.elem_info[elem_id] = (self.current_elem_type, offset + idx)
    
    def _parse_translate(self, line):
        try:
            values = list(map(float, re.split(r'[,\s]+', line)))
            if len(values) >= 3:
                return np.array(values[:3], dtype=np.float32)
        except ValueError:
            pass
        return None
    
    def _parse_rotate(self, line):
        try:
            values = list(map(float, re.split(r'[,\s]+', line)))
            if len(values) >= 7:
                # 前6个值是旋转轴的两个点 (x1,y1,z1, x2,y2,z2)
                # 第7个值是旋转角度(度)
                axis_point1 = np.array(values[:3], dtype=np.float32)
                axis_point2 = np.array(values[3:6], dtype=np.float32)
                angle = np.float32(values[6])  
                return (axis_point1, axis_point2, angle)
        except ValueError:
            pass
        return None


# 使用示例
if __name__ == "__main__":
    parser = AbaqusInpParser(R"example.inp")
    assembly = parser.parse()
    
     # 打印解析结果示例
    print(f"Assembly contains {len(assembly.instances)} instances")
    for instance in assembly.instances:
        part = instance.ref_part
        print(f"\nInstance '{instance.name}' references part '{part.name}'")
        print(f"  Translation: {instance.translate}")
        print(f"  Rotation: {instance.rotation}")
        
        print(f"  Part has {len(part.node_array)} nodes and {sum(len(e) for e in part.elem_types.values())} elements")
        print(f"  Element types: {list(part.elem_types.keys())}")
        
        # 示例：访问特定单元的信息
        if part.elem_info:
            elem_id = next(iter(part.elem_info))  # 获取第一个单元ID
            elem_type, elem_idx = part.elem_info[elem_id]
            nodes = part.elem_types[elem_type][elem_idx]
            print(f"\nExample element {elem_id}:")
            print(f"  Type: {elem_type}")
            print(f"  Node indices: {nodes}")
            print(f"  Node coordinates:")
            for node_id in nodes:
                node_idx = part.node_id_to_index[node_id]
                print(f"    {node_id}: {part.node_array[node_idx]}")